from locust import FastHttpUser, task

# root@server:/home/yangxianpku/code-gathering/LLMs/deploy/vllm# locust -f locust_client.py 
# http://36.111.148.95:8089/?tab=charts
# locust -f locust_qwen_7b.py --headless -u 20 -r 2 --run-time 10m

class ChatUser(FastHttpUser):
    host = "http://127.0.0.1:1025" 
    @task
    def chat(self):
        self.client.post(url="/v1/chat/completions", 
                        headers={
                            "Content-Type"  : "application/json",
                            "Authorization" : "" 
                        },
                        json={
                            "model": "qwen_7b",
                            "messages": [
                                {
                                    "role": "system",
                                    "content": "你是一个旅游达人."
                                },
                                {
                                    "role": "user",
                                    "content": "推荐几个四川省的特色旅游景点并给出推荐原因！<|im_start|>assistant\n"
                                }
                            ]
                        }
                    )